#ifndef QUANT_MATMUL_V2_BLOCK_H
#define QUANT_MATMUL_V2_BLOCK_H

#include "quant_matmul_v2_base.h"

namespace AscendC{

class QuantMatmulV2BaseBlock{

public:
  __aicore__ inline QuantMatmulV2BaseBlock(){}
  __aicore__ inline void init(const QuantMatmulV2TilingData *tiling_data);
  __aicore__ inline void updateBlockCnt(uint64_t m_tile_index, uint64_t n_tile_index);
  __aicore__ inline void initFirstTileBlockIndex();
  __aicore__ inline void initBlockIndex();
  __aicore__ inline void updateBlockIndex();

public:
  QBmmBaseBlockArgs params_;
  bool index_init_ = false;
  uint32_t row_order_;
  uint64_t m_cnt_;
  uint64_t n_cnt_;
  uint64_t m_cnt_tail_;
  uint64_t n_cnt_tail_;
  uint64_t round_;
  uint64_t real_round_;
  uint64_t pre_core_num_;
  uint64_t block_idx_start_;
  uint64_t block_idx_end_;
  const TCubeTiling *matmul_tiling_data_;
};

__aicore__ inline void QuantMatmulV2BaseBlock::init(const QuantMatmulV2TilingData *tiling_data){
    matmul_tiling_data_ = &(tiling_data->matmul_tiling);
    const L2cacheParams &tiling_l2 = tiling_data->tile_l2cache_tiling;
    params_.m_tile_cnt_l2 = static_cast<uint64_t>(tiling_l2.m_tile_cnt_l2);
    params_.n_tile_cnt_l2 = static_cast<uint64_t>(tiling_l2.n_tile_cnt_l2);
    params_.m_total_cnt = (static_cast<uint64_t>(matmul_tiling_data_->M) + matmul_tiling_data_->singleCoreM - 1) / matmul_tiling_data_->singleCoreM;
    params_.n_total_cnt = (static_cast<uint64_t>(matmul_tiling_data_->N) + matmul_tiling_data_->singleCoreN - 1) / matmul_tiling_data_->singleCoreN;
    m_cnt_ = tiling_l2.m_tile_block;
    n_cnt_ = tiling_l2.n_tile_block;
    params_.total_tile_cnt = m_cnt_ * n_cnt_;
    m_cnt_tail_ = params_.m_total_cnt - (params_.m_tile_cnt_l2 - 1) * m_cnt_;
    n_cnt_tail_ = params_.n_total_cnt - (params_.n_tile_cnt_l2 - 1) * n_cnt_;
    round_ = (params_.total_tile_cnt + matmul_tiling_data_->usedCoreNum - 1) / matmul_tiling_data_->usedCoreNum;
    pre_core_num_ = params_.total_tile_cnt % matmul_tiling_data_->usedCoreNum;
    params_.m_cnt_use = m_cnt_;
    params_.n_cnt_use = n_cnt_;
    row_order_ = tiling_l2.cal_order > 0 ? tiling_l2.cal_order : ROW_FIRST;
}

__aicore__ inline void QuantMatmulV2BaseBlock::updateBlockCnt(uint64_t m_tile_index, uint64_t n_tile_index){
    params_.m_tile_addr_offset = m_tile_index * m_cnt_ * matmul_tiling_data_->singleCoreM;
    params_.n_tile_addr_offset = n_tile_index * n_cnt_ * matmul_tiling_data_->singleCoreN;


    if ((m_tile_index == (params_.m_tile_cnt_l2 - 1)) && (n_tile_index == (params_.n_tile_cnt_l2 - 1))) {
        params_.total_tile_cnt = m_cnt_tail_ * n_cnt_tail_;
        params_.m_cnt_use = m_cnt_tail_;
        params_.n_cnt_use = n_cnt_tail_;
    } else if (m_tile_index == (params_.m_tile_cnt_l2 - 1)) {
        params_.total_tile_cnt = m_cnt_tail_ * n_cnt_;
        params_.m_cnt_use = m_cnt_tail_;
        params_.n_cnt_use = n_cnt_;
    } else if (n_tile_index == (params_.n_tile_cnt_l2 - 1)) {
        params_.total_tile_cnt = m_cnt_ * n_cnt_tail_;
        params_.m_cnt_use = m_cnt_;
        params_.n_cnt_use = n_cnt_tail_;
    } else {
        params_.total_tile_cnt = m_cnt_ * n_cnt_tail_;
        params_.m_cnt_use = m_cnt_;
        params_.n_cnt_use = n_cnt_;
    }

    round_ = DequantBmm::ceil_div(params_.total_tile_cnt, static_cast<uint64_t>(matmul_tiling_data_->usedCoreNum));
    pre_core_num_ = params_.total_tile_cnt - (round_ - 1) * matmul_tiling_data_->usedCoreNum;
    if (pre_core_num_ == 0) {
        pre_core_num_ = static_cast<uint64_t>(matmul_tiling_data_->usedCoreNum);
    }
}

__aicore__ inline void QuantMatmulV2BaseBlock::initFirstTileBlockIndex(){
    params_.m_tile_addr_offset = 0;
    params_.n_tile_addr_offset = 0;
    index_init_ = true;
    block_idx_end_ = pre_core_num_;
    if (pre_core_num_ == 0){
        pre_core_num_ = matmul_tiling_data_->usedCoreNum;
    }
    uint64_t pre_total_block = 0U;
    if (block_idx < pre_core_num_) {
        // row_order == 1
        params_.index = block_idx * round_;
        real_round_ = round_;
    }else {
        params_.index = block_idx * (round_ - 1) + pre_core_num_;
        real_round_ = round_ - 1;
    }
}

__aicore__ inline void QuantMatmulV2BaseBlock::initBlockIndex(){
    if (!index_init_) {
        block_idx_start_ = block_idx_end_;
    } else {
        block_idx_start_ = 0;
        index_init_ = true;
    }

    block_idx_end_ = (block_idx_start_ + pre_core_num_) % matmul_tiling_data_ -> usedCoreNum;
    uint64_t index_start = block_idx_start_;
    uint64_t index_end = block_idx_end_;

    if (index_start < index_end){
        if (block_idx < index_start){
            params_.index = block_idx * (round_ - 1);
            real_round_ =  round_ - 1;
        } else if (block_idx < index_end){
            params_.index = index_start * (round_ - 1) + (block_idx - index_start) * round_;
            real_round_ = round_;
        } else {
            params_.index = (index_start * (round_ - 1) + pre_core_num_ * round_ + (block_idx - index_end) * (round_ - 1));
            real_round_ = round_ - 1;
        }
    } else if (index_end < index_start){
        if (block_idx < index_end){
            params_.index = block_idx * round_;
            real_round_ = round_;
        } else if (block_idx < index_start){
            params_.index = index_end * round_ + (block_idx - index_end) * (round_ - 1);
            real_round_ = round_ - 1;
        } else {
            params_.index = (index_end * round_ + (index_start - index_end) * (round_ - 1) + (block_idx - index_start) * round_);
            real_round_ = round_;
        }

    } else {
        params_.index = block_idx * round_;
        real_round_ = round_;
    }
}

__aicore__ inline void QuantMatmulV2BaseBlock::updateBlockIndex(){
    params_.index += 1;
}


} // namespace AscendC

#endif // QUANT_MATMUL_V2_BLOCK_H